import os
import json
from collections import defaultdict

class UnionFind:
    def __init__(self):
        self.parent = {}
    
    def find(self, u):
        if u not in self.parent:
            self.parent[u] = u
        while self.parent[u] != u:
            self.parent[u] = self.parent[self.parent[u]]  # 路径压缩
            u = self.parent[u]
        return u
    
    def union(self, u, v):
        root_u = self.find(u)
        root_v = self.find(v)
        if root_u != root_v:
            self.parent[root_v] = root_u

def parse_obj(obj_path):
    vertices = []
    tex_coords = []
    normals = []
    faces = []
    other_lines = []

    with open(obj_path, 'r') as f:
        for line in f:
            stripped = line.strip()
            if not stripped:
                continue
            if stripped.startswith('v '):
                parts = stripped.split()
                vertices.append(list(map(float, parts[1:4])))
            elif stripped.startswith('vt'):
                parts = stripped.split()
                tex_coords.append(list(map(float, parts[1:])))
            elif stripped.startswith('vn'):
                parts = stripped.split()
                normals.append(list(map(float, parts[1:])))
            elif stripped.startswith('f'):
                parts = stripped.split()
                face = []
                for part in parts[1:]:
                    indices = part.split('/')
                    v_idx = int(indices[0]) - 1
                    vt_idx = int(indices[1]) - 1 if len(indices) > 1 and indices[1] else None
                    vn_idx = int(indices[2]) - 1 if len(indices) > 2 and indices[2] else None
                    face.append((v_idx, vt_idx, vn_idx))
                faces.append(face)
            else:
                other_lines.append(line.rstrip())
    return vertices, tex_coords, normals, faces, other_lines

def split_obj(obj_path, output_prefix="output"):
    vertices, tex_coords, normals, faces, other_lines = parse_obj(obj_path)
    
    uf = UnionFind()
    for face in faces:
        if face:
            first_v = face[0][0]
            for vertex in face[1:]:
                uf.union(first_v, vertex[0])
    
    groups = defaultdict(list)
    for v_idx in range(len(vertices)):
        root = uf.find(v_idx)
        groups[root].append(v_idx)
    
    metadata = {
        "num_original_vertices": len(vertices),
        "num_original_tex_coords": len(tex_coords),
        "num_original_normals": len(normals),
        "groups": []
    }
    
    for group_idx, (root, v_indices) in enumerate(groups.items(), 1):
        group_v = set(v_indices)
        group_faces = [face for face in faces if face and face[0][0] in group_v]
        
        vt_indices = set()
        vn_indices = set()
        for face in group_faces:
            for vertex in face:
                if vertex[1] is not None:
                    vt_indices.add(vertex[1])
                if vertex[2] is not None:
                    vn_indices.add(vertex[2])
        
        # 创建顶点映射
        new_vertices = []
        v_original_indices = []
        for idx in range(len(vertices)):
            if idx in group_v:
                v_original_indices.append(idx)
                new_vertices.append(vertices[idx])
        
        # 创建纹理坐标映射
        new_tex_coords = []
        vt_original_indices = []
        for idx in range(len(tex_coords)):
            if idx in vt_indices:
                vt_original_indices.append(idx)
                new_tex_coords.append(tex_coords[idx])
        
        # 创建法线映射
        new_normals = []
        vn_original_indices = []
        for idx in range(len(normals)):
            if idx in vn_indices:
                vn_original_indices.append(idx)
                new_normals.append(normals[idx])
        
        # 构建输出内容
        output = []
        output.extend(other_lines)
        
        for v in new_vertices:
            output.append(f"v {' '.join(map(str, v))}")
        for vt in new_tex_coords:
            output.append(f"vt {' '.join(map(str, vt))}")
        for vn in new_normals:
            output.append(f"vn {' '.join(map(str, vn))}")
        
        for face in group_faces:
            face_line = []
            for vertex in face:
                parts = []
                parts.append(str(v_original_indices.index(vertex[0]) + 1))  # 当前组中的新索引
                vt_val = vt_original_indices.index(vertex[1]) + 1 if vertex[1] is not None else None
                vn_val = vn_original_indices.index(vertex[2]) + 1 if vertex[2] is not None else None
                
                if vt_val or vn_val:
                    parts.append(str(vt_val) if vt_val is not None else '')
                    if vn_val:
                        parts.append(str(vn_val))
                face_part = '/'.join(parts)
                face_line.append(face_part)
            output.append(f"f {' '.join(face_line)}")
        
        output_path = f"{output_prefix}_{group_idx}.obj"
        with open(output_path, 'w') as f:
            f.write('\n'.join(output))
        print(f"Created: {output_path}")
        
        # 保存元数据信息
        group_metadata = {
            "output_file": output_path,
            "v_map": v_original_indices,
            "vt_map": vt_original_indices,
            "vn_map": vn_original_indices
        }
        metadata["groups"].append(group_metadata)
    
    # 保存元数据到JSON文件
    metadata_path = f"{output_prefix}_metadata.json"
    with open(metadata_path, 'w') as f:
        json.dump(metadata, f, indent=4)
    print(f"Metadata saved to {metadata_path}")

if __name__ == "__main__":
    input_obj = "qiaoduntai.obj"
    split_obj(input_obj)
'''
我的目的不是优化这段代码，目前我收到了新的需求，
拆分出来的单个obj会有实际测出来的ply点云文件和它做配准，
然后点云ply和单个obj配准之后，点云需要依靠拆分时候保存的元信息变换到整体的对应位置去，
因此有两件事需要你完成，一方面修改这份代码，保存更多的元信息，以实现点云和单个obj配准之后能够变换到整体去，
另一方面实现配准和利用元信息变换回整体的代码，也就是说第二份代码根据保存的元信息把若干拆分的obj和对应测量点云读取，
并把测量点云放到整体的指定位置并保存为新的ply
'''
